#include "timer.hpp"
#include "magic_enum.hpp"
#include "exceptions.hpp"
#include "timing.hpp"

namespace hms {

template<typename SolverType>
SolverBase<SolverType>::SolverBase(){
	m_constants.fill( signalCode::uninitialised );
	m_settings.useSourceFields.fill( false );
}

template<typename SolverType>
SolverBase<SolverType>::SolverBase(
	std::array<scalar, nConstants>&& constants,
	std::array<Field<MeshType>, nFields>&& fields,
	std::array<ConstantFieldVar<MeshType>, nConstantFields>&& constantFields,
	std::array<SourceFieldVar<MeshType>, nSourceFields>&& sourceFields
) :
	m_constants      { std::move(constants     ) },
	m_fields         { std::move(fields        ) },
	m_constantFields { std::move(constantFields) },
	m_sourceFields   { std::move(sourceFields  ) }
{}


template<typename SolverType>
void SolverBase<SolverType>::assertInitialisedMembers() const {
	#ifndef NDEBUG
	for ( scalar c : m_constants )
		assert( c != signalCode::uninitialised &&
		"Constant is uninitialised!" );
	
	auto assertNotEmpty = [](const auto& field){
		assert( !field.empty() && "Field is empty!" );
	};
	for ( const Field<MeshType>& field : m_fields )
		assertNotEmpty(field);

	for ( const ConstantFieldVar<MeshType>& field : m_constantFields )
		std::visit(assertNotEmpty, field);

	visitUsedSourceFields( assertNotEmpty );

	#endif
}

template<typename SolverType>
void SolverBase<SolverType>::assertSameMesh() const {

	#ifndef NDEBUG
	const MeshType* meshPtr { &( this->m_fields[0].mesh() ) };

	auto test = [&](const auto& field ){
		assert( &( field.mesh() ) == meshPtr &&
		"Fields do not point to the same mesh!" );
	};
	for ( const Field<MeshType>& field : m_fields )
		test(field);
	for (const ConstantFieldVar<MeshType>& fieldVar : m_constantFields)
		std::visit(test, fieldVar);
	
	visitUsedSourceFields(test);

	#endif
}

template<typename SolverType>
void SolverBase<SolverType>::checkFieldSizes() const {
	this->derived().checkFieldSizesImp();
}

template<typename SolverType>
void SolverBase<SolverType>::validate() const {

	logger().debug("Validating solver...");

	logger().debug("Checking initialised members...");
	this->assertInitialisedMembers();
	logger().debug("Done.");

	logger().debug("Checking initialised solver...");
	assert( this->isInitialised() && "Solver not initialised!" );
	logger().debug("Done.");

	logger().debug("Checking fields pointing to the same mesh...");
	this->assertSameMesh();
	logger().debug("Done.");

	logger().debug("Checking fields having the correct size...");
	this->checkFieldSizes();
	logger().debug("Done.");

	/* this may be difficult to check during import, so we use exceptions
	 * rather than asserts to make it work in release mode */
	logger().debug("Checking for existing boundary conditions...");
	int i {0};
	for ( const Field<MeshType>& field : m_fields ){
		std::string fieldNameStr {
			magic_enum::enum_name( static_cast<FieldName>(i) )
		};
		logger().debug( "Checking field \"{}\"...", fieldNameStr );
		if ( !field.checkBoundaryConditions() )
			throw InputError(
				"Field \"" + fieldNameStr
				+ "\" has missing boundary condition(s)!"
			);
		logger().debug("Done.");
		++i;
	}
	logger().debug("Done.");
	
	logger().debug("Requesting derived solver to perform validation...");
	this->derived().validateImp();
	logger().debug("Derived solver validation completed.");
}

template<typename SolverType>
void SolverBase<SolverType>::initialise(){
	logger().debug("Initialising solver...");

	this->derived().initialiseImp();
	m_isInitialised = true;

	validate();
}

template<typename SolverType>
bool SolverBase<SolverType>::isInitialised() const {
	return m_isInitialised;
}

template<typename SolverType>
scalar SolverBase<SolverType>::calculateTimeStep() const {
	return this->derived().calculateTimeStepImp();
}

template<typename SolverType>
void SolverBase<SolverType>::calculateFluxes(){
	this->derived().calculateFluxesImp();
}

template<typename SolverType>
void SolverBase<SolverType>::updateSourceFields(){
	visitUsedSourceFields( [](auto& sf){ sf.updateValues(); } );
}

template<typename SolverType>
void SolverBase<SolverType>::calculateSources(){
	/* The implementation should check for used source fields */
	this->derived().calculateSourcesImp();
}

template<typename SolverType>
void SolverBase<SolverType>::calculateNewState(){
	this->derived().calculateNewStateImp();
}

template<typename SolverType>
void SolverBase<SolverType>::updateBoundaries(){
	this->derived().updateBoundariesImp();
}

template<typename SolverType>
void SolverBase<SolverType>::advance( Time& time ){
	this->time(time);

	assert( !time.atEnd() &&
		"End time already reached when calling SolverBase::advance!");
	assert( isInitialised() &&
		"Solver is not initialised!");

	Timer t;

	#define LOG_DEBUG(...) \
		logger().debug( __VA_ARGS__ ); \
		setTiming(m_execTimings, LoopStep::output, t)

	LOG_DEBUG( "At time step #{}, time: {:.5f}s, time step length: {:.5f}s",
		time.nSteps(), time.current(), time.step()
	);

	LOG_DEBUG("Updating boundary fields...");
	updateBoundaries();
	setTiming(m_execTimings, LoopStep::boundary, t);

	LOG_DEBUG("Updating source terms...");
	updateSourceFields();
	setTiming(m_execTimings, LoopStep::updateSources, t);

	LOG_DEBUG("Computing fluxes...");
	calculateFluxes();
	setTiming(m_execTimings, LoopStep::fluxes, t);

	LOG_DEBUG("Adding source terms...");
	calculateSources();
	setTiming(m_execTimings, LoopStep::sources, t);

	LOG_DEBUG("Computing state variables of new time step...");
	calculateNewState();
	setTiming(m_execTimings, LoopStep::newState, t);

	LOG_DEBUG("Advancing time...");
	time.advance();
	setTiming(m_execTimings, LoopStep::other, t);
	LOG_DEBUG("New time: {:.5f}s", time.current() );

	if ( time.atEnd() )
		return;

	time.step( calculateTimeStep() );
	setTiming(m_execTimings, LoopStep::timestep, t);
	LOG_DEBUG( "New time step length: {:.5f}s", time.step() );

	#undef LOG_DEBUG
}


template<typename SolverType>
auto SolverBase<SolverType>::time() const -> const Time& {
	assert( m_time && "Time is not set!" );
	return *m_time;
}

template<typename SolverType>
void SolverBase<SolverType>::time( const Time& t ){
	m_time = std::experimental::make_observer( &t );
}

template<typename SolverType>
auto SolverBase<SolverType>::mesh() const -> const MeshType& {
	// assertSameMesh();
	return m_fields[0].mesh();
}


template<typename SolverType>
auto SolverBase<SolverType>::constants() const ->
	const std::array<scalar, nConstants>&
{ return m_constants; }

template<typename SolverType>
auto SolverBase<SolverType>::constants() ->
	std::array<scalar, nConstants>&
{ return m_constants; }

template<typename SolverType>
auto SolverBase<SolverType>::constant( ConstantName name ) -> scalar& {
	return m_constants[ static_cast<int>(name) ];
}

template<typename SolverType>
auto SolverBase<SolverType>::constant( ConstantName name ) const -> scalar {
	return m_constants[ static_cast<int>(name) ];
}

template<typename SolverType>
auto SolverBase<SolverType>::fields() const ->
	const std::array<Field<MeshType>, nFields>&
{ return m_fields; }

template<typename SolverType>
auto SolverBase<SolverType>::fields() ->
	std::array<Field<MeshType>, nFields>&
{ return m_fields; }

template<typename SolverType>
auto SolverBase<SolverType>::field( FieldName name ) const ->
	const Field<MeshType>&
{
	return m_fields[ static_cast<int>(name) ];
}

template<typename SolverType>
auto SolverBase<SolverType>::field( FieldName name ) -> Field<MeshType>& {
	return m_fields[ static_cast<int>(name) ];
}


template<typename SolverType>
constexpr auto SolverBase<SolverType>::zFieldName() -> FieldName {
	return SolverType::zFieldNameImp();
}

template<typename SolverType>
auto SolverBase<SolverType>::zField() const -> const Field<MeshType>& {
	return this->field( zFieldName() );
}

template<typename SolverType>
auto SolverBase<SolverType>::zField() -> Field<MeshType>& {
	return this->field( zFieldName() );
}


template<typename SolverType>
auto SolverBase<SolverType>::constantFields() const ->
	const std::array<ConstantFieldVar<MeshType>, nConstantFields>&
{ return m_constantFields; }

template<typename SolverType>
auto SolverBase<SolverType>::constantFields() ->
	std::array<ConstantFieldVar<MeshType>, nConstantFields>&
{ return m_constantFields; }

template<typename SolverType>
auto SolverBase<SolverType>::constantField( ConstantFieldName name ) -> ConstantFieldVar<MeshType>&
{
	return m_constantFields[ static_cast<int>(name) ];
}

template<typename SolverType>
auto SolverBase<SolverType>::constantField( ConstantFieldName name ) const ->
	const ConstantFieldVar<MeshType>&
{
	return m_constantFields[ static_cast<int>(name) ];
}


template<typename SolverType>
auto SolverBase<SolverType>::sourceFields() const ->
	const std::array<SourceFieldVar<MeshType>, nSourceFields>&
{ return m_sourceFields; }

template<typename SolverType>
auto SolverBase<SolverType>::sourceFields() ->
	std::array<SourceFieldVar<MeshType>, nSourceFields>&
{ return m_sourceFields; }

template<typename SolverType>
auto SolverBase<SolverType>::sourceField( SourceFieldName name ) -> 
	SourceFieldVar<MeshType>&
{
	return m_sourceFields[ static_cast<int>(name) ];
}

template<typename SolverType>
auto SolverBase<SolverType>::sourceField( SourceFieldName name ) const
	-> const SourceFieldVar<MeshType>&
{
	return m_sourceFields[ static_cast<int>(name) ];
}


template<typename SolverType>
constexpr bool SolverBase<SolverType>::useSourceField(SourceFieldName s) const {
	return settings().useSourceFields[static_cast<int>(s)];
}

template<typename SolverType>
void SolverBase<SolverType>::useSourceField(SourceFieldName s, bool b){
	settings().useSourceFields[static_cast<int>(s)] = b;
}

namespace internal {

template<typename T, typename Visitor>
void visitUsedSourceFields( T&& solver, Visitor&& v ){
	/* no copies */
	static_assert( std::is_reference_v<decltype(solver)> );
	using U = remove_qualifiers<T>;
	using MeshType = typename U::MeshType;
	for ( int i{0}; i<U::nSourceFields; ++i){
		auto name { static_cast<typename U::SourceFieldName>(i) };
		if ( solver.useSourceField(name) ){
			transcribe_const_t<T, SourceFieldVar<MeshType>>& sourceVar {
				solver.sourceField(name)
			};
			std::visit( std::forward<Visitor>(v), sourceVar);
		}
	}
}

} // namespace internal

template<typename SolverType>
template<typename Visitor>
void SolverBase<SolverType>::visitUsedSourceFields( Visitor&& v ) const {
	internal::visitUsedSourceFields( *this, std::forward<Visitor>(v) );
}

template<typename SolverType>
template<typename Visitor>
void SolverBase<SolverType>::visitUsedSourceFields( Visitor&& v ){
	internal::visitUsedSourceFields( *this, std::forward<Visitor>(v) );
}


template<typename SolverType>
auto SolverBase<SolverType>::storage() const -> const Field<MeshType>&
{ return m_storage; }

template<typename SolverType>
auto SolverBase<SolverType>::storage() -> Field<MeshType>&
{ return m_storage; }


#ifdef SOLVER_STORE_FLUXES
	
template<typename SolverType>
auto SolverBase<SolverType>::fluxes() const -> const Field<MeshType>& {
	return m_fluxes;
}

template<typename SolverType>
auto SolverBase<SolverType>::fluxes() -> Field<MeshType>& {
	return m_fluxes;
}

#endif

template<typename SolverType>
auto SolverBase<SolverType>::settings() const -> const SettingsType& {
	return m_settings;
}

template<typename SolverType>
auto SolverBase<SolverType>::settings() -> SettingsType& {
	return m_settings;
}

template<typename SolverType>
auto SolverBase<SolverType>::execTimings() const -> const ExecTimings& {
	return m_execTimings;
}

template<typename SolverType>
auto SolverBase<SolverType>::logger() const -> log::logger& {
	return mesh().logger();
}

} // end namespace hms